# -*- coding: utf-8 -*-

from yacs.config import CfgNode as CN


CONFIG = CN()

CONFIG.PHASE = "train"
CONFIG.DEVICE = "cuda:0"

# -----------------------------------------------------------------------------
# INPUT
# -----------------------------------------------------------------------------
CONFIG.INPUT = CN()
# Values to be used for image normalization
CONFIG.INPUT.PIXEL_MEAN = (0.4914, 0.4822, 0.4465)
# Values to be used for image normalization
CONFIG.INPUT.PIXEL_STD = (0.2023, 0.1994, 0.2010)

# -----------------------------------------------------------------------------
# Transform
# -----------------------------------------------------------------------------
CONFIG.TRANSFORM = CN()
CONFIG.TRANSFORM.CROP_SIZE = 32
CONFIG.TRANSFORM.CROP_PADDING_SIZE = 4

# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
CONFIG.DATASET = CN()
CONFIG.DATASET.DATA_ROOT_FOLDER = r"G:\public_dataset\cifar"
CONFIG.DATASET.NUM_CLASS = 10

# -----------------------------------------------------------------------------
# DataLoader
# -----------------------------------------------------------------------------
CONFIG.DATALOADER = CN()
CONFIG.DATALOADER.BATCH_SIZE = 128
CONFIG.DATALOADER.NUM_WORKERS = 4

# -----------------------------------------------------------------------------
# Train
# -----------------------------------------------------------------------------

CONFIG.TRAIN = CN()
CONFIG.TRAIN.EPOCHS = 120
CONFIG.TRAIN.START_EPOCH = 1
CONFIG.TRAIN.RESUME = False
CONFIG.TRAIN.SAVED_MODEL = ""

# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
CONFIG.MODEL = CN()
CONFIG.MODEL.NAME = "squeezenet"

# -----------------------------------------------------------------------------
# Optimizer
# -----------------------------------------------------------------------------
CONFIG.OPTIMIZER = CN()
CONFIG.OPTIMIZER.LEARNING_RATE = 0.1
CONFIG.OPTIMIZER.WEIGHT_DECAY = 5e-4
CONFIG.OPTIMIZER.MILESTONES = (30, 60, 90, )
